package com.hujie.sqlqueryes.service;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import org.apache.http.HttpHost;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.concurrent.TimeUnit;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalListener;
import com.google.common.cache.RemovalNotification;

/**
 * Elasticsearch连接管理器
 * 用于管理多个用户的独立连接
 */
@Service
public class ElasticsearchConnectionManager {

    private static final Logger logger = LoggerFactory.getLogger(ElasticsearchConnectionManager.class);

    // 使用Guava Cache存储用户ID与其对应Elasticsearch连接的映射关系
    // 设置5分钟不使用自动过期
    private final Cache<String, ElasticsearchClient> clientCache;
    private final Cache<String, RestClient> restClientCache;

    public ElasticsearchConnectionManager() {
        // 创建移除监听器，用于在缓存项被移除时关闭连接
        RemovalListener<String, ElasticsearchClient> clientRemovalListener = 
            (RemovalNotification<String, ElasticsearchClient> notification) -> {
                if (notification.getValue() != null) {
                    logger.info("用户 {} 的Elasticsearch客户端连接已移除，原因: {}", 
                               notification.getKey(), notification.getCause());
                }
            };
            
        RemovalListener<String, RestClient> restClientRemovalListener = 
            (RemovalNotification<String, RestClient> notification) -> {
                if (notification.getValue() != null) {
                    try {
                        notification.getValue().close();
                        logger.info("用户 {} 的RestClient连接已关闭，原因: {}", 
                                   notification.getKey(), notification.getCause());
                    } catch (IOException e) {
                        logger.warn("关闭用户 {} 的RestClient连接时出错: {}", 
                                   notification.getKey(), e.getMessage());
                    }
                }
            };

        // 初始化缓存，设置5分钟不使用自动过期
        clientCache = CacheBuilder.newBuilder()
                .expireAfterAccess(5, TimeUnit.MINUTES)
                .removalListener(clientRemovalListener)
                .build();
                
        restClientCache = CacheBuilder.newBuilder()
                .expireAfterAccess(5, TimeUnit.MINUTES)
                .removalListener(restClientRemovalListener)
                .build();
    }

    /**
     * 为指定用户创建或更新Elasticsearch连接
     *
     * @param userId 用户唯一标识符
     * @param host   Elasticsearch主机地址
     * @param port   Elasticsearch端口号
     * @return 是否连接成功
     */
    public boolean connectUser(String userId, String host, int port) {
        logger.info("用户 {} 尝试连接到Elasticsearch {}:{}", userId, host, port);

        // 关闭该用户之前的连接（如果存在）
        disconnectUser(userId);

        // 创建新的连接
        RestClient restClient = null;
        try {
            restClient = RestClient.builder(new HttpHost(host, port)).build();
            ElasticsearchTransport transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);

            // 测试连接
            client.info();

            // 存储连接到缓存
            restClientCache.put(userId, restClient);
            clientCache.put(userId, client);

            logger.info("用户 {} 成功连接到Elasticsearch", userId);
            return true;
        } catch (Exception e) {
            // 如果创建连接过程中出现异常，确保关闭已创建的资源
            if (restClient != null) {
                try {
                    restClient.close();
                } catch (Exception closeException) {
                    logger.warn("关闭用户 {} 的Elasticsearch连接时出错: {}", userId, closeException.getMessage());
                }
            }
            logger.error("用户 {} 连接Elasticsearch失败: {}", userId, e.getMessage(), e);
            return false;
        }
    }

    /**
     * 断开指定用户的Elasticsearch连接
     *
     * @param userId 用户唯一标识符
     */
    public void disconnectUser(String userId) {
        try {
            RestClient restClient = restClientCache.getIfPresent(userId);
            if (restClient != null) {
                restClient.close();
                restClientCache.invalidate(userId);
                clientCache.invalidate(userId);
                logger.info("用户 {} 的Elasticsearch连接已断开", userId);
            }
        } catch (Exception e) {
            logger.error("关闭用户 {} 的Elasticsearch连接时出错: {}", userId, e.getMessage(), e);
        }
    }

    /**
     * 获取指定用户的Elasticsearch客户端
     *
     * @param userId 用户唯一标识符
     * @return Elasticsearch客户端，如果不存在则返回null
     */
    public ElasticsearchClient getClient(String userId) {
        return clientCache.getIfPresent(userId);
    }

    /**
     * 获取当前用户的Elasticsearch客户端
     * 从UserContext中获取当前用户ID
     *
     * @return Elasticsearch客户端，如果不存在则返回null
     */
    public ElasticsearchClient getCurrentClient() {
        String userId = UserContext.getUserId();
        if (userId == null) {
            return null;
        }
        return getClient(userId);
    }

    /**
     * 检查指定用户是否已连接
     *
     * @param userId 用户唯一标识符
     * @return 是否已连接
     */
    public boolean isConnected(String userId) {
        ElasticsearchClient client = clientCache.getIfPresent(userId);
        if (client == null) {
            return false;
        }

        try {
            client.info();
            return true;
        } catch (Exception e) {
            logger.warn("用户 {} 的Elasticsearch连接已失效: {}", userId, e.getMessage());
            // 从缓存中移除失效的连接
            disconnectUser(userId);
            return false;
        }
    }

    /**
     * 检查当前用户是否已连接
     * 从UserContext中获取当前用户ID
     *
     * @return 是否已连接
     */
    public boolean isCurrentConnected() {
        String userId = UserContext.getUserId();
        if (userId == null) {
            return false;
        }
        return isConnected(userId);
    }
}